package cn.glutnn.handler;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class ChatWebSocketHandler extends TextWebSocketHandler {
    private static final ConcurrentHashMap<Integer, WebSocketSession> sessions = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        // 连接建立时不做任何操作，等待认证消息
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        JsonNode jsonNode = new ObjectMapper().readTree(message.getPayload());

        // 处理已读状态消息
        if (jsonNode.has("type") && "read_status".equals(jsonNode.get("type").asText())) {
            Integer senderId = jsonNode.get("senderId").asInt();
            Integer receiverId = jsonNode.get("receiverId").asInt();

            // 转发已读状态给原消息发送方
            WebSocketSession senderSession = sessions.get(receiverId);
            if (senderSession != null && senderSession.isOpen()) {
                senderSession.sendMessage(message);
            }
            return;
        }

        // 处理认证消息
        if (jsonNode.has("type") && "auth".equals(jsonNode.get("type").asText())) {
            Integer userId = jsonNode.get("userId").asInt();
            session.getAttributes().put("userId", userId);
            sessions.put(userId, session);
            return;
        }

        // 处理心跳消息
        if (jsonNode.has("type") && "ping".equals(jsonNode.get("type").asText())) {
            session.sendMessage(new TextMessage("{\"type\":\"pong\"}"));
            return;
        }

        // 处理普通消息
        Integer senderId = jsonNode.get("senderId").asInt();
        Integer receiverId = jsonNode.get("receiverId").asInt();
        String messageContent = jsonNode.get("message").asText();

        // 转发消息给接收者
        WebSocketSession receiverSession = sessions.get(receiverId);
        if (receiverSession != null && receiverSession.isOpen()) {
            receiverSession.sendMessage(new TextMessage(message.getPayload()));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        Integer userId = (Integer) session.getAttributes().get("userId");
        if (userId != null) {
            sessions.remove(userId);
        }
    }

    public void sendMessage(Integer receiverId, String message) throws IOException {
        WebSocketSession session = sessions.get(receiverId);
        if (session != null && session.isOpen()) {
            // 构造一个标准的消息格式
            Map<String, Object> messageData = new HashMap<>();
            messageData.put("type", "message");
            messageData.put("message", message);

            ObjectMapper mapper = new ObjectMapper();
            String jsonMessage = mapper.writeValueAsString(messageData);

            session.sendMessage(new TextMessage(jsonMessage));
        }
    }
}
